import numpy as np
import pandas as pd

def get_inflation_forecasts(df, melt_id_vars):
    """Calculate inflation forecasts from
    individual-level forecasts of GDP price
    deflator.
    
    Args:
        df (df): Dataframe containing individual forecasts
        melt_id_vars (str or list): Variables for converting
            into long format when calculating dates for which
            the forecast was made
    
    Returns:
        df: Dataframe containing inflation forecasts
    """

    # Calculate annualized quarter-on-quarter inflation rates
    # Here, PGDP1 is forecast for previous quarter; PGDP2, PGDP3, ..., PGDP6
    # are forecasts for quarter t, quarter t + 1, ..., quarter t + 4 from now 
    # (see pp. 20-22 of Survey of Professional Forecasters Documentation);
    # PGDPA and PGDPB annual fourth-quarter over fourth-quarter
    # forecasts with A and B for year t and year t + 1
    # (see p. 20 of Survey of Professional Forecasters Documentation)
    for ii in range(2, 6 + 1):
        df['INF{}'.format(ii)] = 100 * ((df['PGDP{}'.format(ii)] / df['PGDP{}'.format(ii - 1)]) ** 4 - 1)

    # Melt the data into long format
    value_vars_list = ['INF{}'.format(ii) for ii in range(2, 6 + 1)]
    df = pd.melt(df, id_vars = melt_id_vars, value_vars = value_vars_list)

    # Calculate dates *for which* the forecast was made
    df['quarters_ahead'] = df['variable'].map(lambda s : int(s[-1]) - 2)
    df['DATE'] = df['DATE'] + df['quarters_ahead'].apply(lambda x: pd.DateOffset(months = x * 3))

    # Pivot back into wide format
    df = df.pivot_table(index = melt_id_vars, columns = 'variable', values = 'value')
    df = df.reset_index()
    df['DATE'] = df['DATE'].apply(lambda x: str(x.year)) + ':0' + df['DATE'].apply(lambda x: str(x.quarter))

    # Rename variables in accordance to those in SPF
    # forecast error statistics: 
    # SPFfor_Step1 means forecast for current quarter,
    # SPFfor_Step2 -- quarter t + 1, ..., SPFfor_Step5 -- quarter t + 4
    for var_name in df.columns:
        if 'INF' in var_name:
            tt = int(var_name[-1])
            new_var_name = 'SPFfor_Step{}'.format(tt - 1)
            df = df.rename(columns={var_name: new_var_name})

    return df